import torch
import torch.nn as nn
import torch.nn.functional as F


class ResBlock(nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ResBlock, self).__init__()
        self.convs = nn.ModuleList(
            [nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size-1)//2 * d) for d in dilation]
        )

    def forward(self, x):
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1) + x
        return x

class Generator(nn.Module):
    def __init__(self, upsample_rates=(8, 8, 2, 2), upsample_kernel_sizes=(16, 16, 4, 4)):
        super(Generator, self).__init__()
        self.conv_pre = nn.Conv1d(80, 512, 7, 1, padding=3)
        self.ups = nn.ModuleList()
        self.resblocks = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(nn.ConvTranspose1d(512 // (2 ** i), 512 // (2 ** (i + 1)), k, u, padding=(k - u) // 2))
            self.resblocks.append(ResBlock(512 // (2 ** (i + 1))))

        self.conv_post = nn.Conv1d(512 // (2 ** len(upsample_rates)), 1, 7, 1, padding=3)

    def forward(self, x):
        x = F.leaky_relu(self.conv_pre(x), 0.1)
        for up, resblock in zip(self.ups, self.resblocks):
            x = F.leaky_relu(up(x), 0.1)
            x = resblock(x)
        x = torch.tanh(self.conv_post(x))
        return x

class PeriodDiscriminator(nn.Module):
    def __init__(self, period):
        super(PeriodDiscriminator, self).__init__()
        self.period = period
        self.convs = nn.ModuleList([
            nn.Conv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)),
            nn.Conv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
            nn.Conv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
            nn.Conv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
            nn.Conv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
        ])
        self.conv_post = nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))

    def forward(self, x):
        fmap = []
        b, c, t = x.shape
        if t % self.period != 0:
            n_pad = self.period - (t % self.period)
            x = F.pad(x, (0, n_pad), "reflect")
        x = x.view(b, c, x.size(2) // self.period, self.period)
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)
        return x, fmap

class MultiPeriodDiscriminator(nn.Module):
    def __init__(self):
        super(MultiPeriodDiscriminator, self).__init__()
        self.discriminators = nn.ModuleList([
            PeriodDiscriminator(p) for p in [2, 3, 5, 7, 11]
        ])

    def forward(self, x):
        ret = []
        fmaps = []
        for d in self.discriminators:
            out, fmap = d(x)
            ret.append(out)
            fmaps.append(fmap)
        return ret, fmaps

class ScaleDiscriminator(nn.Module):
    def __init__(self):
        super(ScaleDiscriminator, self).__init__()
        self.convs = nn.ModuleList([
            nn.Conv1d(1, 128, 15, 1, padding=7),
            nn.Conv1d(128, 128, 41, 4, groups=4, padding=20),
            nn.Conv1d(128, 256, 41, 4, groups=16, padding=20),
            nn.Conv1d(256, 512, 41, 4, groups=16, padding=20),
            nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20),
            nn.Conv1d(1024, 1024, 5, 1, padding=2)
        ])
        self.conv_post = nn.Conv1d(1024, 1, 3, 1, padding=1)

    def forward(self, x):
        fmap = []
        for conv in self.convs:
            x = F.leaky_relu(conv(x), 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)
        return x, fmap

class MultiScaleDiscriminator(nn.Module):
    def __init__(self):
        super(MultiScaleDiscriminator, self).__init__()
        self.discriminators = nn.ModuleList([ScaleDiscriminator() for _ in range(3)])
        self.meanpools = nn.ModuleList([nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)])

    def forward(self, x):
        ret = []
        fmaps = []
        for i, d in enumerate(self.discriminators):
            if i != 0:
                x = self.meanpools[i-1](x)
            out, fmap = d(x)
            ret.append(out)
            fmaps.append(fmap)
        return ret, fmaps

